import torch
import torch.nn as nn


def torch_custom_conv2d(inputs, output_channel, kernel_size):
    """
    使用PyTorch实现的自定义2D卷积函数 (padding=0, stride=1)

    参数:
        inputs: 输入张量，形状为 (batch_size, C_in, H, W)
        output_channel: 输出通道数
        kernel_size: 卷积核大小（正方形）

    返回:
        outputs: 卷积结果，形状为 (batch_size, output_channel, H_out, W_out)
    """
    # 获取输入维度信息
    batch_size, C_in, H_in, W_in = inputs.shape

    # 计算输出特征图尺寸
    H_out = H_in - kernel_size + 1
    W_out = W_in - kernel_size + 1

    # 初始化卷积核 (output_channel, C_in, kernel_size, kernel_size)
    # 使用Xavier初始化方法
    kernel = nn.init.xavier_uniform_(torch.empty(output_channel, C_in, kernel_size, kernel_size))

    # 初始化输出张量
    outputs = torch.zeros(batch_size, output_channel, H_out, W_out)

    # 执行卷积操作
    for b in range(batch_size):  # 遍历批次
        for c_out in range(output_channel):  # 遍历输出通道
            for i in range(H_out):  # 高度方向滑动
                for j in range(W_out):  # 宽度方向滑动
                    # 提取输入的局部区域 (C_in, kernel_size, kernel_size)
                    input_patch = inputs[b, :, i:i + kernel_size, j:j + kernel_size]

                    # 获取当前输出通道对应的卷积核
                    current_kernel = kernel[c_out]

                    # 计算卷积：元素相乘后求和
                    outputs[b, c_out, i, j] = torch.sum(input_patch * current_kernel)

    return outputs
